// treemodel.cpp
#include "treemodel.h"
#include <QDataStream>
#include <QMimeData>
#include <QDebug>
#include <QIODevice>
#include <QTreeView>

TreeModel::TreeModel(const QStringList &headers, QObject *parent) : QAbstractItemModel(parent)
{
    m_headers = std::move(headers);
    // 初始化根节点
    m_rootNode.children.resize(1); // 确保至少有一个Laminate节点
}

void TreeModel::setTreeData(const Node& rootData)
{
    beginResetModel();
    m_rootNode = rootData;
    endResetModel();
    if (auto view = qobject_cast<QTreeView*>(QObject::parent())) {
        view->expandAll();
    }
    m_isChanged = false;
}

QModelIndex TreeModel::index(int row, int column, const QModelIndex &parent) const
{
    if (!hasIndex(row, column, parent))
        return QModelIndex();

    if (!parent.isValid()) {
        // 根节点下只有Laminate节点
        if (row == 0) {
            return createIndex(row, column, static_cast<quintptr>(0)); // 明确转换为quintptr
        }
        return QModelIndex();
    } else {
        const auto& plyNodes = m_rootNode.children[0].children;
        // Laminate节点的子节点(Ply)
        if (parent.row() == 0) { // 只有Laminate节点有子节点
            if (row >= 0 && row < plyNodes.size()) {
                // 内部指针设为行号+1 (避免0冲突)
                return createIndex(row, column, static_cast<quintptr>(row + 1));
            }
        }
    }
    return QModelIndex();
}

QModelIndex TreeModel::parent(const QModelIndex &index) const
{
    if (!index.isValid())
        return QModelIndex();

    quintptr id = index.internalId();
    if (id == 0) {
        // Laminate节点的父节点是根节点
        return QModelIndex();
    } else {
        // Ply节点的父节点是Laminate节点(第0行)
        return createIndex(0, 0, static_cast<quintptr>(0));
    }
}

int TreeModel::rowCount(const QModelIndex &parent) const
{
    if (!parent.isValid()) {
        // 根节点只有Laminate一行
        return 1;
    } else if (!parent.parent().isValid() && parent.row() == 0) {
        // Laminate节点的子节点数量
        return m_rootNode.children[0].children.size();
    }
    return 0; // Ply节点没有子节点
}

int TreeModel::columnCount(const QModelIndex &) const
{
    return m_headers.size();
}

QVariant TreeModel::data(const QModelIndex &index, int role) const
{
    if (!index.isValid())
        return QVariant();
    auto col = index.column();
    const auto& key = m_headers[col];
    if (!index.parent().isValid() && index.row() == 0) {
        // Laminate节点
        const auto& dataMap = m_rootNode.children[0].data;
        if (role == Qt::DisplayRole || role == Qt::EditRole) {
            return dataMap.value(key);
        }
    } else if (index.internalId() > 0) {
        // Ply节点
        int plyRow = index.internalId() - 1; // 转换回实际行号
        const auto& plyNodes = m_rootNode.children[0].children;
        if (plyRow >= 0 && plyRow < plyNodes.size()) {
            const auto& plyData = plyNodes[plyRow].data;
            switch (role) {
            case Qt::DisplayRole:
            case Qt::EditRole:
                return plyData.value(key);
                break;
            case Qt::BackgroundRole:
                return plyData.value("color");
            case Qt::TextAlignmentRole:
                return Qt::AlignCenter;
            case Qt::UserRole:
                return plyData;
            }
        }
    }
    return QVariant();
}

QVariant TreeModel::headerData(int section, Qt::Orientation orientation, int role) const
{
    if (orientation == Qt::Horizontal && role == Qt::DisplayRole) {
        return m_headers[section];
    }
    return QVariant();
}

Qt::ItemFlags TreeModel::flags(const QModelIndex &index) const
{
    Qt::ItemFlags f = QAbstractItemModel::flags(index);

    if (!index.isValid()) {
        // 根节点：允许放置（用于在末尾添加）
        return f | Qt::ItemIsDropEnabled;
    }

    if (!index.parent().isValid() && index.row() == 0) {
        // Laminate节点：允许放置
        f |= Qt::ItemIsDropEnabled;
    } else if (index.internalId() > 0) {
        // Ply节点：允许拖动
        f |= Qt::ItemIsDragEnabled;
    }
    return f;
}

QStringList TreeModel::mimeTypes() const
{
    return QStringList() << "application/x-ply";
}

QMimeData* TreeModel::mimeData(const QModelIndexList &indexes) const
{
    if (indexes.isEmpty())
        return nullptr;

    QMimeData *mime = new QMimeData;
    QByteArray encoded;
    QDataStream ds(&encoded, QIODevice::WriteOnly);

    QSet<int> rows;
    // 只处理Ply节点
    for (const QModelIndex &idx : indexes) {
        if (idx.isValid() && idx.internalId() > 0) {
            rows.insert(idx.internalId() - 1);
        }
    }
    for (const auto& row : rows)
        ds << row;

    mime->setData("application/x-ply", encoded);
    return mime;
}

bool TreeModel::dropMimeData(const QMimeData *data, Qt::DropAction action,
                             int row, int, const QModelIndex &parent)
{
    if (action != Qt::MoveAction)
        return false;
    if (!data->hasFormat("application/x-ply"))
        return false;

    QByteArray encoded = data->data("application/x-ply");
    QDataStream ds(&encoded, QIODevice::ReadOnly);
    QList<int> sourceRows;
    int r;
    while (!ds.atEnd()) {
        ds >> r;
        sourceRows.append(r);
    }
    if (sourceRows.isEmpty())
        return false;

    auto& plyNodes = m_rootNode.children[0].children;
    // 确定目标位置
    int targetIndex = row;
    if (targetIndex == -1) {
        // 拖到空白处，添加到末尾
        targetIndex = plyNodes.size();
    }
    // 通知模型刷新
    beginResetModel();
    // 收集要移动的节点
    QVector<QVariantMap> movedNodes;
    // 从大到小排序，以便安全移除
    std::sort(sourceRows.begin(), sourceRows.end(), std::greater<int>());
    for (int srcRow : sourceRows) {
        if (srcRow >= 0 && srcRow < plyNodes.size()) {
            movedNodes.prepend(plyNodes.takeAt(srcRow).data);
        }
    }

    // 调整目标位置（考虑已移除的项目）
    for (int srcRow : sourceRows) {
        if (srcRow < targetIndex) {
            targetIndex--;
        }
    }

    // 插入到目标位置
    // for (const auto& node : movedNodes) {
    //     plyNodes.insert(targetIndex, node);
    //     targetIndex++;
    // }

    auto stepInc = m_callback();
    // 计算目标位置的Sequence和Step值
    for (auto& node : movedNodes) {
        // 获取上下行的Sequence和Step值
        QString prevSequence, nextSequence;
        double prevStep = 0, nextStep = 0;

        // 上一行数据
        if (targetIndex > 0) {
            const auto& prevNode = plyNodes[targetIndex - 1].data;
            prevSequence = prevNode["Sequence"].toString();
            prevStep = prevNode["Step"].toDouble();
        }

        // 下一行数据
        if (targetIndex < plyNodes.size()) {
            const auto& nextNode = plyNodes[targetIndex].data;
            nextSequence = nextNode["Sequence"].toString();
            nextStep = nextNode["Step"].toDouble();
        }

        // 应用拖放规则
        if (targetIndex == 0) {
            // 规则4：拖动到Laminate下方
            const auto& lamNode = m_rootNode.children[0].data;
            node["Sequence"] = lamNode["Sequence"];
            node["Step"] = lamNode["Step"];
        } else if (prevSequence != nextSequence /*&& !nextSequence.isEmpty()*/) {
            // 规则1：Sequence不一致
            node["Sequence"] = prevSequence;
            node["Step"] = prevStep + stepInc;
        } else if (prevSequence == nextSequence && prevStep != nextStep) {
            // 规则2：Sequence一致，Step不一致
            node["Sequence"] = prevSequence;
            node["Step"] = (prevStep + nextStep) / 2;
        } else if (prevSequence == nextSequence && prevStep == nextStep) {
            // 规则3：Sequence和Step都一致
            node["Sequence"] = prevSequence;
            node["Step"] = prevStep;
        } else {
            // 默认情况：保持原有值
        }

        plyNodes.insert(targetIndex, node);
        targetIndex++;
    }
    endResetModel();

    if (auto view = qobject_cast<QTreeView*>(QObject::parent())) {
        view->expandAll();
    }
    m_isChanged = true;
    return true;
}

Qt::DropActions TreeModel::supportedDropActions() const
{
    return Qt::MoveAction;
}
